from dataclasses import dataclass
from flask import Blueprint, request
from flask_jwt_extended import create_access_token, create_refresh_token, jwt_required, get_jwt_identity

from pydantic import BaseModel, ValidationError
from typing import Optional

from utils import sm4
from common.constant import message as msg

from do.dto.auth import LoginRequest, RegisterReq, ResetPassReq
from utils import from_result

from utils._db import gen_session
from sqlalchemy import select
from models.users import Users

import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


router = Blueprint("auth", __name__, url_prefix="/auth")


@router.post("/signin")
def auth_login():
    """登录接口"""
    try:
        login_data: Optional[LoginRequest] = LoginRequest(**request.get_json())
        # print(login_data)
    except ValidationError as e:
        return from_result(code=400, message=msg.INVALID_PARAMS)

    ss = gen_session()

    # SM4加密
    gm_sm4: str = sm4.encode_sm4(login_data.password)

    user: Optional[Users] = (
        ss.query(
            Users.username,
            Users.password,
            Users.phone_num,
        )
        .filter(Users.phone_num == login_data.username)
        .first()
    )

    # 每次操作完成后执行关闭操作
    ss.close()

    if not user:
        return from_result(code=401, message=msg.USER_NOT_FOUND)

    if user.password != gm_sm4:
        return from_result(code=401, message=msg.WRONG_PASS)

    access_token = create_access_token(identity=user.phone_num, fresh=True)
    refresh_token = create_refresh_token(identity=user.phone_num)

    return from_result(message=f"Welcome {user.username}!", data={"token": access_token, "refresh_token": refresh_token})


@router.post("/register")
def auth_register():
    """注册接口 TODO"""
    return ""


@router.post("/reset_pass")
@jwt_required(refresh=False)
def auth_reset_password():
    """密码重置"""
    try:
        reset_req: ResetPassReq = ResetPassReq(**request.get_json())
    except ValidationError as e:
        return from_result(code=400, message=msg.INVALID_PARAMS)

    phone_num = get_jwt_identity()

    @dataclass
    class PassCallback:
        password: str

    ss = gen_session()
    # oldPass: PassCallback = Users.select(Users.password).where(Users.phone_num == phone_num).first()
    oldPass: PassCallback = ss.query(Users.password).filter(Users.phone_num == phone_num).first()

    if oldPass.password == sm4.encode_sm4(reset_req.oldPass):

        encrypted_pass: str = sm4.encode_sm4(reset_req.newPass)

        # query: ModelUpdate = Users.update(password=encrypted_pass).where(Users.phone_num == phone_num)
        user: Users = ss.query(Users).filter(Users.phone_num == phone_num).first()
        user.password = encrypted_pass
        try:
            ss.commit()
        except Exception as e:
            logger.warning(e)
            return from_result(code=500, message=msg.PASS_FAILED)

        return from_result(message=msg.OPT_SUCCESS)

    return from_result(code=400, message=msg.WRONG_PASS)


# See https://flask-jwt-extended.readthedocs.io/en/stable/refreshing_tokens.html#explicit-refreshing-with-refresh-tokens
@router.route("/refresh", methods=["POST"])
@jwt_required(refresh=True)
def refresh():
    """刷新fresh token"""
    identity = get_jwt_identity()
    access_token = create_access_token(identity=identity)
    return from_result(message=msg.OPT_SUCCESS, data={"token": access_token})
